/*
 * H264PropDetector.cpp
 *
 *  Created on: 2017年3月7日
 *      Author: chuanjiang.zh
 */

#include "H264PropDetector.h"
#include "H264NaluParser.h"
#include "H265PropParser.h"


namespace av
{

H264PropDetector::H264PropDetector():
		m_codecID(),
		m_prop(),
		m_vps(),
		m_sps(),
		m_pps()
{
}

H264PropDetector::~H264PropDetector()
{
}

H264PropDetector::H264PropDetector(int codecID):
		m_codecID(),
		m_prop(),
		m_vps(),
		m_sps(),
		m_pps()
{
	setCodec(codecID);
}

bool H264PropDetector::setCodec(int codecID)
{
	bool done = false;
	if (codecID == MEDIA_CODEC_H264 || codecID == MEDIA_CODEC_HEVC)
	{
		m_codecID = codecID;
		done = true;
	}
	return done;
}

bool H264PropDetector::detect(const uint8_t* data, int size, bool& isKeyFrame, int& naluType)
{
    if (!H264NaluParser::startWithH264Code(data, size))
    {
        return false;
    }

    if (m_codecID == MEDIA_CODEC_H264)
    {
    	return detectH264(data, size, isKeyFrame, naluType);
    }
    else if (m_codecID == MEDIA_CODEC_HEVC)
    {
    	return detectH265(data, size, isKeyFrame, naluType);
    }

    return false;
}

bool H264PropDetector::isReady() const
{
	return m_prop.size() > 0;
}

std::string H264PropDetector::getProp() const
{
	return m_prop;
}

void H264PropDetector::clear()
{
	m_prop.clear();
	m_vps.clear();
	m_sps.clear();
	m_pps.clear();
}

std::string H264PropDetector::getVps() const
{
	return m_vps;
}

std::string H264PropDetector::getSps() const
{
	return m_sps;
}

std::string H264PropDetector::getPps() const
{
	return m_pps;
}

bool H264PropDetector::detectH264(const uint8_t* data, int size, bool& isKeyFrame, int& naluType)
{
	bool found = false;

	size_t begin = 0;

	while (true)
	{
		size_t pos = H264NaluParser::findH264StartCode(data, size, begin + 3);
		if (pos != -1)
		{
			isKeyFrame |= parseH264Nalu(data + begin, pos - begin, naluType);
			begin = pos;
		}
		else
		{
			break;
		}
	}

	isKeyFrame |= parseH264Nalu(data + begin, size - begin, naluType);

	if (m_sps.size() && m_pps.size())
	{
		std::string prop;
		prop.append((char*)H264NaluParser::s_startCode, H264NaluParser::START_CODE_LENGTH);
		prop.append(m_sps);
		prop.append((char*)H264NaluParser::s_startCode, H264NaluParser::START_CODE_LENGTH);
		prop.append(m_pps);

		m_prop = prop;

		found = true;
	}

	return found;
}

bool H264PropDetector::detectH265(const uint8_t* data, int size, bool& isKeyFrame, int& naluType)
{
	bool found = false;
    size_t begin = 0;

    while (true)
    {
        size_t pos = H264NaluParser::findH264StartCode(data, size, begin + 3);
        if (pos != -1)
        {
        	isKeyFrame |= parseH265Nalu(data + begin, pos - begin, naluType);
            begin = pos;
        }
        else
        {
            break;
        }
    }

    isKeyFrame |= parseH265Nalu(data + begin, size - begin, naluType);

    if (m_sps.size() && m_pps.size())
    {
        std::string prop;
        prop.append((char*)H264NaluParser::s_startCode, H264NaluParser::START_CODE_LENGTH);
        prop.append(m_vps);
        prop.append((char*)H264NaluParser::s_startCode, H264NaluParser::START_CODE_LENGTH);
        prop.append(m_sps);
        prop.append((char*)H264NaluParser::s_startCode, H264NaluParser::START_CODE_LENGTH);
        prop.append(m_pps);

        m_prop = prop;

        found = true;
    }

    return found;
}


bool H264PropDetector::parseH264Nalu(const uint8_t* data, size_t length, int& naluType)
{
	bool isKeyFrame = true;

    NaluPacket nalu;
    if (!H264NaluParser::parseNalu(data, length, nalu))
    {
        return false;
    }

    naluType = nalu.type;

    if (nalu.type == NaluPacket::NALU_SPS)
    {
        m_sps.assign((char*)nalu.data, nalu.length);
    }
    else if (nalu.type == NaluPacket::NALU_PPS)
    {
        m_pps.assign((char*)nalu.data, nalu.length);
    }
    else if (nalu.type == NaluPacket::NALU_IFRAME)
    {
    	//
    }
    else if (nalu.type == NaluPacket::NALU_SEI)
    {
    	//
    }
    else
    {
    	isKeyFrame = false;
    }

    return isKeyFrame;
}

bool H264PropDetector::parseH265Nalu(const uint8_t* data, size_t length, int& naluType)
{
	bool isKeyFrame = true;

    NaluPacket nalu;
    if (!H264NaluParser::parseNalu(data, length, nalu))
    {
        return false;
    }

    nalu.type = (nalu.data[0] & 0x7E) >> 1;

    naluType = nalu.type;

    if (nalu.type == 32)
    {
        m_vps.assign((char*)nalu.data, nalu.length);
    }
    else if (nalu.type == 33)
    {
        m_sps.assign((char*)nalu.data, nalu.length);
    }
    else if (nalu.type == 34)
    {
        m_pps.assign((char*)nalu.data, nalu.length);
    }
    else if (nalu.type == 39)
    {
    	/// sei
    }
    else if (nalu.type == 19 || nalu.type == 20)
    {
    	/// iframe
    }
    else if (nalu.type == 21)
    {
    	///   NAL_UNIT_CODED_SLICE_CRA,       // 21
    }
    else
    {
    	isKeyFrame = false;
    }

    return isKeyFrame;
}



} /* namespace av */
